import utils
import numpy as np
import pandas as pd
import networkx as nx
import matplotlib
import matplotlib.pyplot as plt
import warnings


def visualize(traceID):
    spanMap = utils.fetch_trace_by_id(traceID)
    spans = list(spanMap.values())

    stats = pd.read_csv(f'../data/normal_traces/1/normal_traces_stats.csv').set_index(['service', 'name'])

    for idx, span in enumerate(spans):
        span.idx = idx
        span.duration_anomaly_score_3sigma = (span.duration - stats.loc[(span.service, span.name)]['duration_mean']) \
                                             / stats.loc[(span.service, span.name)]['duration_std']
        span.self_duration_anomaly_score_3sigma = (span.self_duration - stats.loc[(span.service, span.name)]['self_duration_mean']) \
                                                  / stats.loc[(span.service, span.name)]['self_duration_std']

    # 创建有向图
    G = nx.DiGraph()

    # 添加 node
    for i in range(len(spans)):
        G.add_node(i)

    # 添加 edge
    for span in spans:
        if span.parent is not None:
            G.add_edge(span.parent.idx, span.idx)

    # 异常点
    node_colors = np.full(len(spans), 'lightblue')
    node_colors[[span.idx for span in spans if span.duration_anomaly_score_3sigma > 1]] = 'yellow'
    node_colors[[span.idx for span in spans if span.duration_anomaly_score_3sigma > 2]] = 'orange'
    node_colors[[span.idx for span in spans if span.duration_anomaly_score_3sigma > 3]] = 'red'
    node_colors[[span.idx for span in spans if span.self_duration_anomaly_score_3sigma > 3]] = '#8B1D0F'

    # 绘制图形
    warnings.filterwarnings("ignore", category=matplotlib.MatplotlibDeprecationWarning)
    pos = nx.nx_agraph.graphviz_layout(G, prog='dot')
    nx.draw(G, pos, with_labels=True, node_size=100, node_color=node_colors, font_size=12, arrows=True)
    plt.show()

